//---cat rawtcp.c---
// Run as root or SUID 0, just datagram no data/payload
#include <unistd.h>
#include <stdio.h>
#include <sys/socket.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <string.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <errno.h>
#include "libtcp_capture.h"

#include <pthread.h>
// Packet length
#define PCKT_LEN 8192

struct result{
    double avg_conn;
};

struct worker_t{
    pthread_t pid;
    char buffer[PCKT_LEN];
    struct sockaddr_in din;
    int num_conn;
    int run_flag;
    struct result *res;
};

struct handle {
    int tot_td;
    int tot_conn;
    struct worker_t workers[100];
};

// Simple checksum function, may use others such as Cyclic Redundancy Check, CRC
unsigned short csum(unsigned short *buf, int len)
{
    unsigned long sum;
    for(sum=0; len>0; len--)
        sum += *buf++;
    sum = (sum >> 16) + (sum &0xffff);
    sum += (sum >> 16);
    return (unsigned short)(~sum);
}

uint32_t uint16_checksum(uint16_t *data, size_t n)
{
    uint32_t sum=0;
    while (n >= 2) {
        sum += (uint32_t)ntohs(*data);
        data++;        
        n -= 2;
    }                      
    if (n == 1) {
        sum += (uint32_t)ntohs(*(uint8_t *)data);
    }
    return sum;
}       

/*
static uint16_t tcp_checksum(uint16_t *data, size_t n, struct iphdr *ip)
{
    uint32_t sum = uint16_checksum(data, n);
    uint16_t sum2;
    sum += uint16_checksum((uint16_t *)(void *)&ip->saddr,
            sizeof(ip->saddr));
    sum += uint16_checksum((uint16_t *)(void *)&ip->daddr,
            sizeof(ip->daddr));
    sum += ip->protocol + n;
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum = (sum & 0xFFFF) + (sum >> 16);
    sum2 = htons(sum);
    sum2 = ~sum2;
    if (sum2 == 0) {
        return 0xFFFF;
    }
    return sum2;
}
*/

void fill_packet(char *argv[], char *buffer)
{
    struct iphdr *ip = (struct iphdr *)buffer;
    struct tcphdr *tcp = (struct tcphdr *)(buffer + sizeof(struct iphdr));
    char *opt = (char *)tcp + sizeof(struct tcphdr);

    memset(buffer, 0, PCKT_LEN);

    // IP structure
    ip->ihl = 5;
    ip->version = 4;
    ip->tos = 0;
    ip->tot_len = sizeof(struct iphdr) + sizeof(struct tcphdr);
    ip->id = htons(54321);
    ip->frag_off = 0;
    ip->ttl = 64;
    ip->protocol = 6; // TCP
    ip->check = 0; // Done by kernel

    // Source IP, modify as needed, spoofed, we accept through command line argument
    ip->saddr = inet_addr(argv[1]);
    // Destination IP, modify as needed, but here we accept through command line argument
    ip->daddr = inet_addr (argv[3]);

    // The TCP structure. The source port, spoofed, we accept through the command line
    tcp->source = htons(atoi(argv[2]));
    // The destination port, we accept through command line
    tcp->dest = htons(atoi(argv[4]));
    tcp->seq = htonl(1);
    tcp->ack_seq = 0;
    tcp->doff = 0;
    tcp->syn = 0;
    tcp->window = htons(32767);
    tcp->check = 0; // Done by kernel
    tcp->urg_ptr = 0;
    // IP checksum calculation
    ip->check = csum((unsigned short *) buffer, ip->tot_len >> 1);

    opt[0] = 'c';
    opt[1] = 'c';
    opt[2] = 'm';
    opt[3] = 'a';
}

void update_port(char *buffer, unsigned port)
{
    struct tcphdr *tcp = (struct tcphdr *)(buffer + sizeof(struct iphdr));
    tcp->source = htons(port%65530+1);
}

void make_syn(char *buffer)
{
    struct tcphdr *tcp = (struct tcphdr *)(buffer + sizeof(struct iphdr));
    tcp->syn = 1;
    tcp->ack = 0;
    tcp->fin = 0;
    //tcp->check = 0;
    //tcp->check = tcp_checksum((uint16_t *)tcp, sizeof(struct tcphdr), ip);
}

void make_ack(char *buffer)
{
    struct tcphdr *tcp = (struct tcphdr *)(buffer + sizeof(struct iphdr));
    tcp->syn = 0;
    tcp->ack = 1;
    tcp->fin = 0;
}

void make_fin(char *buffer)
{
    struct tcphdr *tcp = (struct tcphdr *)(buffer + sizeof(struct iphdr));
    tcp->syn = 0;
    tcp->ack = 0;
    tcp->fin = 1;
}

double stress_connection_count(int sd, char *buffer, int tot_len, 
                               struct sockaddr_in *din, int total)
{
    // sendto() loop, send every 2 second for 50 counts
    unsigned int count;
    struct timeval start, end;
    unsigned long long start_utime, end_utime, total_utime;
    double result;

    gettimeofday(&start, NULL);
    for(count = 0; count < total; count++) {
        update_port(buffer, count+1);
        /* syn */
        make_syn(buffer);
        if(sendto(sd, buffer, tot_len, MSG_DONTWAIT, (struct sockaddr *)din, sizeof(*din)) < 0) {
            perror("sendto() error");
            return -1;
        }

        /* ack */
        make_ack(buffer);
        if(sendto(sd, buffer, tot_len, MSG_DONTWAIT, (struct sockaddr *)din, sizeof(*din)) < 0) {
            perror("sendto() error");
            return -1;
        }

        /* fin */
        make_fin(buffer);
        if(sendto(sd, buffer, tot_len, MSG_DONTWAIT, (struct sockaddr *)din, sizeof(*din)) < 0) {
            perror("sendto() error");
            return -1;
        }

        /* ack */
        make_ack(buffer);
        if(sendto(sd, buffer, tot_len, MSG_DONTWAIT, (struct sockaddr *)din, sizeof(*din)) < 0) {
            perror("sendto() error");
            return -1;
        }
    }
    gettimeofday(&end, NULL);

    start_utime = start.tv_sec * 1000000 + start.tv_usec;
    end_utime = end.tv_sec * 1000000 + end.tv_usec;
    total_utime = end_utime - start_utime;
    result = ((double)total/(double)total_utime*1000000);
    return result;
}

void *worker_function(void *arg) 
{
    struct worker_t *data = (struct worker_t *)arg;
    struct result *rst = malloc(sizeof(struct result));
    int tot_len = sizeof(struct iphdr) + sizeof(struct tcphdr) + 4;
    int sd = socket(PF_INET, SOCK_RAW, IPPROTO_TCP);
    int one = 1;
    const int *val = &one;

    if(sd < 0) {
        perror("socket() error");
        return NULL;
    }

    /*TODO check malloc */

    // Inform the kernel do not fill up the headers' structure, we fabricated our own
    if(setsockopt(sd, IPPROTO_IP, IP_HDRINCL, val, sizeof(one)) < 0) {
        perror("setsockopt() error");
        return NULL;
    }
    
    while(data->run_flag == 0) {
        usleep(100);
    }

    rst->avg_conn = stress_connection_count(sd, data->buffer, tot_len, &data->din, data->num_conn);
    close(sd);

    return (void *)rst;
}

int create_workers(struct handle *hd)
{
    int i;
    
    for (i = 0; i < hd->tot_td; i++) {
        pthread_create(&hd->workers[i].pid, NULL, worker_function, (void *)&hd->workers[i]);
        printf("create thread [%d]\n", i);
    }
    return 0;
}

int main(int argc, char *argv[])
{
    struct filter_rule_info_user rule;
    struct handle *hd = malloc(sizeof(struct handle));;
    double total_cps;
    int i;
    
    memset(hd, 0, sizeof(struct handle));
    if(argc != 7)
    {
        printf("- Invalid parameters!!!\n");
        printf("- Usage: %s <source hostname/IP> <source port> <target hostname/IP>"
               " <target port> <#of conn> <#of threads>\n", argv[0]);
        exit(-1);
    }

    printf("Using:::::Source IP: %s port: %u, Target IP: %s port: %u.\n", argv[1], atoi(argv[2]), argv[3], atoi(argv[4]));
    
    hd->tot_conn = atoi(argv[5]);
    hd->tot_td = atoi(argv[6]);
    /* fill buffer */
    for (i = 0; i < hd->tot_td; i++) {
        fill_packet(argv, hd->workers[i].buffer);

        hd->workers[i].din.sin_family = AF_INET;
        hd->workers[i].din.sin_port = htons(atoi(argv[4]));
        hd->workers[i].din.sin_addr.s_addr = inet_addr(argv[3]);
 
        hd->workers[i].num_conn = hd->tot_conn/hd->tot_td;
    }

    if (create_workers(hd) < 0) {
        goto end;
    }

    rule.d_ip = inet_addr(argv[3]);
    rule.d_port = htons(atoi(argv[4]));
    if (tcp_capture_add(&rule) < 0) {
        printf("tcp_capture_add failed because %s\n", strerror(errno));
        goto end;
    }
    sleep(1);
    /* run threads */
    printf("threads run...\n");
    for (i = 0; i < hd->tot_td; i++) {
        hd->workers[i].run_flag = 1;
    }

    /* wait all threads */
    for (i = 0; i < hd->tot_td; i++) {
        pthread_join(hd->workers[i].pid, (void **)&hd->workers[i].res);
    }
    
    if (tcp_capture_del(&rule) < 0) {
        printf("tcp_capture_del failed\n");
    }
    /* cal avg result */
    total_cps = 0;
    for (i = 0; i < hd->tot_td; i++) {
        total_cps += hd->workers[i].res->avg_conn;
        free(hd->workers[i].res);
    }
    
    sleep(1);
    printf("total captured packets are: %u packets\n", tcp_capture_get_counter());
    tcp_capture_zero();
 
    printf("threads stoped.\n");
    printf("Avg: %f conn/sec\n", total_cps);

end:
    free(hd);
    return 0;
}

